#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图片提取工具函数
提供兼容性更好的图片提取方法
"""

import re
from typing import Dict, List, Tuple, Optional
from docx import Document
from docx.oxml.ns import qn


def extract_images_from_docx_robust(docx_path: str) -> Dict[str, bytes]:
    """
    健壮的图片提取方法，支持多种docx版本
    
    Args:
        docx_path: docx文件路径
        
    Returns:
        Dict[str, bytes]: {关系ID: 图片二进制数据}
    """
    try:
        doc = Document(docx_path)
        
        # 获取所有图片关系
        image_parts = {}
        for rel in doc.part.rels.values():
            if "image" in rel.target_ref:
                image_parts[rel.rId] = rel.target_part
        
        return {rid: part.blob for rid, part in image_parts.items()}
        
    except Exception as e:
        print(f"提取图片失败: {e}")
        return {}


def find_image_references_in_paragraph(paragraph, image_parts: dict) -> List[str]:
    """
    在段落中查找图片引用，使用多种兼容方法
    
    Args:
        paragraph: docx段落对象
        image_parts: 图片部件字典
        
    Returns:
        List[str]: 找到的图片关系ID列表
    """
    found_refs = []
    
    for run in paragraph.runs:
        # 方法1: 遍历所有元素查找blip
        try:
            for element in run.element.iter():
                if element.tag.endswith('}blip'):
                    embed_id = element.get(qn('r:embed'))
                    if embed_id and embed_id in image_parts:
                        if embed_id not in found_refs:
                            found_refs.append(embed_id)
        except Exception:
            pass
        
        # 方法2: 在XML中搜索关系ID
        try:
            xml_content = str(run.element.xml)
            for rel_id in image_parts.keys():
                if rel_id in xml_content:
                    if rel_id not in found_refs:
                        found_refs.append(rel_id)
        except Exception:
            pass
        
        # 方法3: 查找特定的图片元素标签
        try:
            # 查找drawing元素
            drawing_tags = [
                '{http://schemas.openxmlformats.org/wordprocessingml/2006/main}drawing',
                '{http://schemas.openxmlformats.org/wordprocessingml/2006/main}pict'
            ]
            
            for tag in drawing_tags:
                for drawing in run.element.findall('.//' + tag):
                    # 在drawing元素中查找embed属性
                    for elem in drawing.iter():
                        embed_id = elem.get(qn('r:embed'))
                        if embed_id and embed_id in image_parts:
                            if embed_id not in found_refs:
                                found_refs.append(embed_id)
        except Exception:
            pass
    
    return found_refs


def get_image_extension_from_content_type(content_type: str) -> str:
    """
    根据内容类型获取图片扩展名
    
    Args:
        content_type: MIME类型
        
    Returns:
        str: 文件扩展名
    """
    if 'jpeg' in content_type.lower() or 'jpg' in content_type.lower():
        return '.jpg'
    elif 'png' in content_type.lower():
        return '.png'
    elif 'gif' in content_type.lower():
        return '.gif'
    elif 'bmp' in content_type.lower():
        return '.bmp'
    elif 'tiff' in content_type.lower():
        return '.tiff'
    elif 'webp' in content_type.lower():
        return '.webp'
    else:
        return '.jpg'  # 默认扩展名


def extract_images_with_context_robust(docx_path: str) -> List[Tuple[str, str, bytes]]:
    """
    健壮的带上下文图片提取
    
    Args:
        docx_path: docx文件路径
        
    Returns:
        List[Tuple[str, str, bytes]]: [(标题, 文件名, 图片数据), ...]
    """
    try:
        doc = Document(docx_path)
        
        # 获取图片关系
        image_parts = {}
        for rel in doc.part.rels.values():
            if "image" in rel.target_ref:
                image_parts[rel.rId] = rel.target_part
        
        if not image_parts:
            return []
        
        results = []
        current_title = "未分类"
        title_counters = {}
        
        # 遍历所有段落
        for paragraph in doc.paragraphs:
            # 检查是否是标题
            if is_heading_paragraph(paragraph):
                current_title = clean_title_text(paragraph.text.strip())
            
            # 查找段落中的图片
            image_refs = find_image_references_in_paragraph(paragraph, image_parts)
            
            for ref_id in image_refs:
                if ref_id in image_parts:
                    image_part = image_parts[ref_id]
                    
                    # 生成文件名
                    if current_title not in title_counters:
                        title_counters[current_title] = 0
                    title_counters[current_title] += 1
                    
                    ext = get_image_extension_from_content_type(image_part.content_type)
                    filename = f"{current_title}-{title_counters[current_title]:02d}{ext}"
                    
                    results.append((current_title, filename, image_part.blob))
        
        return results
        
    except Exception as e:
        print(f"提取图片失败: {e}")
        return []


def is_heading_paragraph(paragraph) -> bool:
    """
    判断段落是否为标题
    
    Args:
        paragraph: docx段落对象
        
    Returns:
        bool: 是否为标题
    """
    try:
        # 检查样式名称
        if paragraph.style and paragraph.style.name:
            style_name = paragraph.style.name.lower()
            if any(keyword in style_name for keyword in ['heading', '标题', 'title']):
                return True
        
        # 检查字体属性
        if paragraph.runs:
            first_run = paragraph.runs[0]
            try:
                if first_run.font.size and first_run.font.size.pt > 12:
                    return True
                if first_run.font.bold:
                    return True
            except:
                pass
        
        # 检查文本模式
        text = paragraph.text.strip()
        if re.match(r'^\d+[\.\s]', text) or re.match(r'^[一二三四五六七八九十]+[\.\s]', text):
            return True
        
        return False
        
    except Exception:
        return False


def clean_title_text(title: str) -> str:
    """
    清理标题文本

    Args:
        title: 原始标题文本

    Returns:
        str: 清理后的标题
    """
    if not title:
        return "未命名标题"

    # 移除extract_前缀（包含UUID）
    title = re.sub(r'^extract_[a-f0-9\-]+[-_]', '', title)

    # 移除数字编号
    title = re.sub(r'^\d+[\.\s]*', '', title)
    title = re.sub(r'^[一二三四五六七八九十]+[\.\s]*', '', title)

    # 移除特殊字符
    title = re.sub(r'[<>:"/\\|?*]', '', title)

    # 限制长度
    if len(title) > 50:
        title = title[:50]

    return title.strip() if title.strip() else "未命名标题"


def test_extraction_methods(docx_path: str):
    """
    测试不同的提取方法
    
    Args:
        docx_path: docx文件路径
    """
    print(f"测试文件: {docx_path}")
    
    # 方法1: 基础提取
    images = extract_images_from_docx_robust(docx_path)
    print(f"方法1 - 基础提取: 找到 {len(images)} 张图片")
    
    # 方法2: 带上下文提取
    context_images = extract_images_with_context_robust(docx_path)
    print(f"方法2 - 上下文提取: 找到 {len(context_images)} 张图片")
    
    for i, (title, filename, image_data) in enumerate(context_images, 1):
        print(f"  {i}: {title} -> {filename} ({len(image_data)} 字节)")


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        test_extraction_methods(sys.argv[1])
    else:
        print("用法: python image_extractor_utils.py <docx文件路径>")
